package com.fxkj.cloud.drms.infra.ws;

import com.fxkj.cloud.dbs.common.model.ws.FxWsMsg;
import com.fxkj.cloud.drms.domain.ws.service.FxWsClient;
import com.fxkj.cloud.drms.infra.persistence.cache.OrganizationCacheRepository;
import com.fxkj.common.controller.BaseController;
import com.fxkj.common.exception.FxException;
import com.fxkj.common.utils.JacksonUtil;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.MoreExecutors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * <p>description: </p>
 * @author ww
 * @date 2025/3/17 15:50
 */
@Component
@Slf4j
public class FxWsClientImpl extends BaseController implements FxWsClient, WebSocketHandler {

    private static final ConcurrentHashMap<UUID, Set<WebSocketSession>> sessions = new ConcurrentHashMap<>();
    private static final AtomicInteger onlineCount = new AtomicInteger(0);
    private final OrganizationCacheRepository organizationCacheRepository;
    // 新增缓存结构
    private static final ConcurrentHashMap<UUID, LinkedBlockingQueue<FxWsMsg>> messageCache = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<UUID, AtomicInteger> messageCounters = new ConcurrentHashMap<>();
    private static final ScheduledExecutorService scheduler = new ScheduledThreadPoolExecutor(
            Runtime.getRuntime().availableProcessors() * 2,
            new ThreadFactory() {
                private final AtomicInteger threadNumber = new AtomicInteger(1);

                @Override
                public Thread newThread(Runnable r) {
                    Thread t = new Thread(r, "ws-scheduler-" + threadNumber.getAndIncrement());
                    t.setDaemon(true); // 设置为守护线程
                    return t;
                }
            },
            new ThreadPoolExecutor.CallerRunsPolicy() // 拒绝策略
    );
    private final WsConfig wsConfig;

    public FxWsClientImpl(OrganizationCacheRepository organizationCacheRepository, WsConfig wsConfig) {
        this.organizationCacheRepository = organizationCacheRepository;
        this.wsConfig = wsConfig;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        UUID userId = null;
        try {
            userId = getUserId();
        } catch (FxException e) {
            throw new RuntimeException(e);
        }

        if (null == userId) {
            log.debug("[websocket] -| websocket开启失败,参数为空");
            return;
        }
        Set<WebSocketSession> existedSessions = sessions.getOrDefault(userId, new HashSet<>());
        existedSessions.add(session);
        sessions.put(userId, existedSessions);
        log.debug("[websocket] -| websocket开启成功");
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        log.info("[websocket] -| websocket收到消息:{}", message);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.debug("[websocket] -| websocket异常,exception:{}", null != exception ? exception.getMessage() : null);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        // 新增连接关闭时的缓存清理 todo 获取用户id
//        UUID userId = getUserId();
//        if (userId != null) {
//            messageCache.remove(userId);
//            messageCounters.remove(userId);
//        }
        session.close();
        log.debug("[websocket] -| websocket关闭,sessionId:{}", session.getId());

    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    @Override
    public void pushNotificationToUser(UUID userId, FxWsMsg wsMsg) {
//        Set<WebSocketSession> webSocketSessions = sessions.getOrDefault(userId, new HashSet<>());
//        if (null == webSocketSessions || webSocketSessions.isEmpty()) {
//            log.debug("[websocket] -| 用户不在线,userId:{}", userId);
//            return;
//        }
//        Boolean someOneClose = false;
//        List<WebSocketSession> removeSessions = new ArrayList<>();
//        for (WebSocketSession webSocketSession : webSocketSessions) {
//            if (webSocketSession.isOpen()) {
//                try {
//                    log.debug("[websocket] -| websocket推送消息,userId:{}", userId);
//                    webSocketSession.sendMessage(new TextMessage(wsMsg.toJsonString()));
//                } catch (IOException e) {
//                    throw new RuntimeException(e);
//                }
//            } else {
//                removeSessions.add(webSocketSession);
//                someOneClose = true;
//                log.debug("[websocket] -| websocket已关闭,无法推送消息,userId:{}", userId);
//            }
//        }
//        webSocketSessions.removeAll(removeSessions);
//        //如果有用户关闭了连接，需要更新缓存
//        if (someOneClose) {
//            sessions.put(userId, webSocketSessions);
//        }
        // 将新消息加入缓存队列
        messageCache.computeIfAbsent(userId, k -> new LinkedBlockingQueue<>()).offer(wsMsg);
        AtomicInteger counter = messageCounters.computeIfAbsent(userId, k -> new AtomicInteger(0));

        // 当达到1000条时立即发送
        if (counter.incrementAndGet() >= wsConfig.getBatchSize()) {
            flushMessages(userId);
            return;
        }

        // 首次添加消息时启动定时任务
        if (counter.get() == 1) {
            scheduler.schedule(() -> flushMessages(userId), wsConfig.getFlushIntervalMs(), TimeUnit.MILLISECONDS);
        }
    }

    @Override
    public void pushNotificationToUsers(Set<UUID> userIds, FxWsMsg wsMsg) {
        Futures.submit(() -> {
            for (UUID userId : userIds) {
                pushNotificationToUser(userId, wsMsg);
            }
        }, MoreExecutors.directExecutor());
    }

    @Override
    public void pushNotificationToOrganization(UUID organizationId, FxWsMsg msg) {
        Set<UUID> notifyUserIds = organizationCacheRepository.getUserIdsByOrgId(organizationId);
        if (null == notifyUserIds) {
            return;
        }
        for (UUID notifyUserId : notifyUserIds) {
            pushNotificationToUser(notifyUserId, msg);
        }
    }

    // 消息批量发送方法
    private void flushMessages(UUID userId) {
        LinkedBlockingQueue<FxWsMsg> queue = messageCache.get(userId);
        if (queue == null || queue.isEmpty()) return;

        List<FxWsMsg> batch = new ArrayList<>();
        queue.drainTo(batch);

        Set<WebSocketSession> sessions = this.sessions.get(userId);
        List<WebSocketSession> removeSessions = new ArrayList<>();
        if (sessions != null) {
            try {
                TextMessage textMessage = new TextMessage(JacksonUtil.writeValueAsString(batch));
                for (WebSocketSession session : sessions) {
                    if (session.isOpen()) {
                        session.sendMessage(textMessage);
                    } else {
                        removeSessions.add(session);
                    }
                }
            } catch (IOException e) {
                log.error("消息批量发送失败", e);
            }
        }
        if (!removeSessions.isEmpty()) {
            sessions.removeAll(removeSessions);
            if (sessions.isEmpty()) {
                this.sessions.remove(userId);
            } else {
                this.sessions.put(userId, sessions);
            }
        }

        messageCounters.get(userId).set(0);
    }


}
